package immutablecollections;

import immutablecollections.exceptions.*;
import immutablecollections.functions.*;

/**
 *
 */

public class ImRoseTreeShapes
{

    private static Function1<ImRoseTree<String>> zWithNodes = new Function1<ImRoseTree<String>>()
    {
        @SuppressWarnings("unchecked")
        @Override
        public ImRoseTree<String> invoke(Object nodes)
        {
            return ImRoseTree.withNodes("z", (ImShelf<ImRoseTree<String>>) nodes);
        }
    };

    // A relabel function that I want to use in a map
    public static Function1<ImRoseTree<String>> relabelTreeFn = FnFactory.on(ImRoseTreeShapes.class).getFnStatic(
            ImRoseTree.class, "relabelTree", ImRoseTree.class);

    // A relabel function that I want to use in a map
    public static Function2<ImRoseTree<String>> relabelTreeUsingFactoryFn = //
    FnFactory.on(ImRoseTreeShapes.class).getFnStatic( //
            ImRoseTree.class, "relabelTreeUsingFactory", LabelFactory.class, ImRoseTree.class);

    public static ImRoseTree<String> relabelTree(ImRoseTree<String> treeToRelabel)
    {
        return relabelTreeUsingFactory(new LabelFactory(), treeToRelabel);
    }

    public static ImRoseTree<String> relabelTreeUsingFactory(LabelFactory labelFactory, ImRoseTree<String> treeToRelabel)
    {
        String myLabel = labelFactory.getLabel();
        return ImRoseTree.withNodes(myLabel, relabelForest(treeToRelabel.getSubTrees(), labelFactory));
    }

    private static ImShelf<ImRoseTree<String>> relabelForest(ImShelf<ImRoseTree<String>> trees,
            LabelFactory labelFactory)
    {
        return trees.map(relabelTreeUsingFactoryFn.invoke(labelFactory));
    }

    public static class LabelFactory
    {
        /**
         *  the labels at each level
         *  a1
         *  |
         *  ------------
         *  b1 b2     b3
         *     |      |
         *     -----  -----
         *     c1 c2  c3 c4
         */
        private int count = 0;

        public String getLabel()
        {
            return "" + (++count);
        }

    }

    public ImList<ImRoseTree<String>> allTreesWithSize(int treeSize)
    {
        ImIndexOutOfBoundsException.check(treeSize, 1, Integer.MAX_VALUE, "treeSize");
        return allTreesWithSize$(treeSize).map(relabelTreeFn);
    }

    private ImList<ImRoseTree<String>> allTreesWithSize$(int n)
    {
        if (n == 1)
            return ImList.on(ImRoseTree.leaf("z"));
        else
            return allForestsWithSize(n - 1).<ImRoseTree<String>> map(zWithNodes);
    }

    private ImList<ImShelf<ImRoseTree<String>>> allForestsWithSize(int n)
    {

        if (n == 0)
        {
            return ImList.on(ImShelf.<ImRoseTree<String>> empty());
        }

        ImList<ImShelf<ImRoseTree<String>>> results = ImList.empty();

        /**
         * Consider the first tree of the forest. It can have size 1 to n. If it has size n then it must be the only
         * tree in the forest
         * 
         *    firstTree
         *    |
         *    --
         *    t1   t2 t3 .. tn
         *         -----------
         *             |
         *         otherTrees
         */

        // For a first tree of each possible size
        for (int i = 1; i <= n; i++)
        {
            // So - the first tree has size i.
            // The remaining trees of the forest will be a forest that has a size of n - i

            // Let's get all possible such forests
            ImList<ImShelf<ImRoseTree<String>>> otherTrees = allForestsWithSize(n - i);

            // For all possible trees of size i
            for (ImRoseTree<String> firstTree : allTreesWithSize$(i))
            {
                // We add the first tree to the front of each of these forests
                ImList<ImShelf<ImRoseTree<String>>> firstTreePlusOtherTrees = cons(firstTree, otherTrees);

                // And add this forest to the results
                results = firstTreePlusOtherTrees.addingAll(results);
            }
        }

        return results;
    }

    private ImList<ImShelf<ImRoseTree<String>>> cons(ImRoseTree<String> tree,
            ImList<ImShelf<ImRoseTree<String>>> allForests)
    {
        ImList<ImShelf<ImRoseTree<String>>> results = ImList.empty();

        for (ImShelf<ImRoseTree<String>> forest : allForests)
        {
            results = ImList.cons(forest.adding(1, tree), results);
        }

        return results.reversed();
    }
}
